#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>

#include "user_threads.h"
#include "user_mutex.h"
#include "user_io.h"

#include "sieve.h"
#include "queue.h"
#include "errno.h"
#include "assert.h"

struct work {
	unsigned long prime_no;
	unsigned long prime;
	uthread_id worker;
};

struct writer_params {
	const char * output_file;
};

struct reader_params {
	const char * input_file;
};

static uthread_mutex_t input_lock;
static uthread_mutex_t output_lock;
static uthread_mutex_t thread_lock;
static struct queue * input;
static struct queue * output;

static int writer_fd;
static int dispatch_fd;
static int workers;

/* the threads */
int dispatch(void * arg);
int worker(void * arg);
int writer(void * arg);
int reader(void * arg);

int dispatch(void * arg)
{
	int fds[2] = {0,};
	char buf[10];
	uthread_id w = -1;
	if (uthread_pipe(fds)) {
		perror("Dispatch pipe");
		uthread_cancel(0);
	}
	dispatch_fd = fds[1];
	for (;;) {
		if (!uthread_read(fds[0], buf, 1)) {
			perror("Dispatch read");
			uthread_cancel(0);
		}
		switch (buf[0]) {
		case 'a':
			/* new input */
			
			if (uthread_mutex_lock(thread_lock)
						!= MUTEX_SUCCESS) {
				write(2, "Dispatch thread mutex lock\n",27);
				uthread_cancel(0);
			}
			w = uthread_create(worker, NULL, UTHREAD_MIN_PRIO);
			if (UTHREAD_FAIL == w) {
				write(2, "Dispatch spawn error\n",21);
				assert(false);
				uthread_cancel(0);
			}
			w = UTHREAD_FAIL;
			workers++;
			if (uthread_mutex_unlock(thread_lock)
						!= MUTEX_SUCCESS) {
				write(2, "Dispatch thread mutex unlock\n",29);
				uthread_cancel(0);
			}
			break;
		case 'q':
			/* quit */
			uthread_exit(0);
			break;
		default:
			/* unexpected */
			uthread_exit(1);
			break;
		}
		uthread_yield();
	}
	close(fds[0]);
	close(fds[1]);
}

int worker(void * arg)
{
	struct work * work = NULL;
	unsigned long * prime_array;
	if (uthread_self() == UTHREAD_MAX_THREADS - 1) {
		if (uthread_mutex_lock(thread_lock) != MUTEX_SUCCESS) {
			write(2, "Worker thread mutex lock\n",25);
			return -1;
		}
	}
	if (uthread_mutex_lock(input_lock) != MUTEX_SUCCESS) {
		write(2, "Worker mutex lock\n",18);
		return -1;
	}
	assert(!queue_is_empty(input));
	if (queue_remove(input, (void **) &work))
		assert(false);
	if (uthread_mutex_unlock(input_lock) != MUTEX_SUCCESS) {
		write(2, "Worker mutex unlock\n",20);
		return -1;
	}
	if (!(prime_array = malloc(work->prime_no * sizeof(*prime_array)))) {
		perror("Worker malloc");
		return -1;
	}
	sieve(prime_array, work->prime_no);
	work->prime = prime_array[work->prime_no - 1];
	work->worker = uthread_self();
	free(prime_array);
	prime_array = NULL;
	if (uthread_mutex_lock(output_lock) != MUTEX_SUCCESS) {
		write(2, "Worker mutex lock\n",18);
		return -1;
	}
	if (queue_add(output, work))
		assert(false);
	if (uthread_mutex_unlock(output_lock) != MUTEX_SUCCESS) {
		write(2, "Worker mutex unlock\n",20);
		return -1;
	}
	if (1 != uthread_write(writer_fd, "b", 1)) {
		write(2,"Worker write\n",13);
	}
	workers--;
	return 0;
}

int writer(void * arg)
{
	struct writer_params * params = arg;
	int worker_ret = -1;
	int out = -1;
	int fds[2] = {0};
	char buf[50];
	int size = -1;
	struct work * work = NULL;
	out = uthread_open(params->output_file, O_WRONLY | O_TRUNC | O_CREAT);
	if (out < 0) {
		perror("Writer open");
		uthread_cancel(0);
	}
	if (uthread_pipe(fds)) {
		perror("Write pipe");
		uthread_cancel(0);
	}
	writer_fd = fds[1];
	for (;;) {
		if (!uthread_read(fds[0], buf, 1)) {
			perror("Writer read");
			uthread_cancel(0);
		}
		switch (buf[0]) {
		case 'b':
			/* new output */
			if (uthread_mutex_lock(output_lock) != MUTEX_SUCCESS) {
				write(2,"Worker mutex lock\n",18);
				uthread_cancel(0);
			}
			if (queue_remove(output, (void **) &work))
				assert(false);
			if (uthread_mutex_unlock(output_lock) != MUTEX_SUCCESS) {
				write(2, "Worker mutex unlock\n",20);
				return -1;
			}
			size = snprintf(buf, 50, "%lu: %lu\n",
					work->prime_no, work->prime);
			if (size <= 0) {
				perror("Writer snprintf");
				uthread_cancel(0);
			}
			uthread_join(work->worker, &worker_ret);
			assert(worker_ret == 0);
			free(work);
			work = NULL;
			uthread_write(out, buf, size);
			size = -1;
			break;
		case 'q':
			/* quit */
			uthread_exit(0);
			break;
		default:
			/* unexpected */
			uthread_exit(1);
			break;
		}
	}
	close(out);
	close(fds[0]);
	close(fds[1]);
	return 0;
}

int reader(void * arg)
{
	struct reader_params * params = arg;
	char mainbuf[1024], auxbuf[1024];
	char *buf = mainbuf, *token = NULL;
	struct work * work = NULL;
	int pos = 0, aux_len = -1;
	int ret = 0;
	int fd = -1;

	fd = uthread_open(params->input_file, O_RDONLY);
	assert(fd > 0);

	while (((ret = uthread_read(fd, mainbuf + pos, 1023 - pos)) > 0) ||
			(errno == EINTR)) {
		aux_len = ret;
		while(mainbuf[aux_len-1] != '\n')
			--aux_len;
		memcpy(auxbuf, mainbuf, aux_len);
		auxbuf[aux_len] = '\0';
		for (buf = auxbuf; ; buf = NULL) {
			token = strtok(buf, "\n");
			if (!token)
				break;
			if (!(work = malloc(sizeof(*work)))) {
				perror("Reader malloc");
				uthread_cancel(0);
			}
			work->prime_no = strtoul(token, NULL, 10);
//			write(1,"P\n",2);
			if (uthread_mutex_lock(input_lock) != MUTEX_SUCCESS) {
				write(2,"Reader mutex lock\n",17);
				uthread_cancel(0);
			}
//			write(1,"I\n",2);
			if (queue_add(input, work))
				assert(false);
//			write(1,"R\n",2);
			if (uthread_mutex_unlock(input_lock) != MUTEX_SUCCESS) {
				write(1, "Reader mutex unlock\n",20);
				uthread_cancel(0);
			}
//			write (1,"Q\n",2);
			if (1 != uthread_write(dispatch_fd, "a", 1)) {
				perror("Reader write");
				uthread_cancel(0);
			}
			work = NULL;
		}
		errno = 0;
		assert(errno != EINTR);
		pos = ret - aux_len;
		memmove(mainbuf, mainbuf + aux_len, pos+1);
	}
	close(fd);
	return 0;
}

int main(int argc, char ** argv)
{
	uthread_id thread[3];
	struct writer_params w;
	struct reader_params r;
	assert(argc == 3);
	r.input_file = argv[1];
	w.output_file = argv[2];
	uthread_init();
	uthread_mutex_init(&thread_lock);
	uthread_mutex_init(&input_lock);
	uthread_mutex_init(&output_lock);
	input = queue_init();
	output = queue_init();
	thread[0] = uthread_create(dispatch, NULL, UTHREAD_MIN_PRIO);
	thread[1] = uthread_create(writer, &w, UTHREAD_MIN_PRIO);
	thread[2] = uthread_create(reader, &r, UTHREAD_MIN_PRIO);
	uthread_join(thread[2], NULL);
	while (workers != 0)
		uthread_yield();
	uthread_write(dispatch_fd, "q", 1);
	uthread_join(thread[0], NULL);
	uthread_write(writer_fd, "q", 1);
	uthread_join(thread[1], NULL);
	uthread_mutex_destroy(thread_lock);
	uthread_mutex_destroy(input_lock);
	uthread_mutex_destroy(output_lock);
	return 0;
}
